import os
import sys
import time

from PIL import Image
from selenium import webdriver

"""
i use librewolf and its for me
from selenium.webdriver.firefox.firefox_binary import FirefoxBinary
browser = webdriver.Firefox(executable_path=os.path.join(os.path.abspath(os.path.dirname(__file__)),'geckodriver.exe'),firefox_binary=FirefoxBinary('C:\\Program Files\\LibreWolf\\librewolf.exe'))
"""


if os.path.exists('tasks')==False:
    os.mkdir('tasks')

browser=input("Which browser will you use/Какой браузер вы будете использовать:\n1.Firefox\n2.Google Chrome\n3.PhantomJS\n? ")
try:
    if browser=="Firefox":
        from selenium.webdriver.firefox.options import Options
        browser = webdriver.Firefox(executable_path=os.path.join(os.path.abspath(os.path.dirname(__file__)),'geckodriver.exe'),firefox_options=Options().add_argument("--headless"))
    elif browser=="Google Chrome":
        from selenium.webdriver.chrome.options import Options
        
        browser = webdriver.Firefox(executable_path=os.path.join(os.path.abspath(os.path.dirname(__file__)),'chromedriver.exe'),chrome_options=Options().add_argument("--headless"))
    elif browser=="PhantomJS":
        browser = webdriver.PhantomJS(executable_path=os.path.join(os.path.abspath(os.path.dirname(__file__)),'phantomjs.exe'))
    else:
        print("Nothing selected, exiting...")
        time.sleep(3)
        sys.exit()
except:
    if browser=="Firefox":
        from selenium.webdriver.firefox.options import Options
        browser = webdriver.Firefox(executable_path=os.path.join(os.path.abspath(os.path.dirname(__file__)),'geckodriver'),firefox_options=Options().add_argument("--headless"))
    elif browser=="Google Chrome":
        from selenium.webdriver.chrome.options import Options
        
        browser = webdriver.Firefox(executable_path=os.path.join(os.path.abspath(os.path.dirname(__file__)),'chromedriver'),chrome_options=Options().add_argument("--headless"))
    elif browser=="PhantomJS":
        browser = webdriver.PhantomJS(executable_path=os.path.join(os.path.abspath(os.path.dirname(__file__)),'phantomjs'))
    else:
        print("Nothing selected, exiting...")
        time.sleep(3)
        sys.exit()
else:
    print("Cant find driver, exiting...")
    time.sleep(3)
    sys.exit()


async def watermark_with_transparency(input_image_path,
                                output_image_path,
                                watermark_image_path,
                                ):
    base_image = Image.open(input_image_path)
    watermark = Image.open(os.path.join(os.path.abspath(os.path.dirname(__file__)),watermark_image_path))
    width1,height1=watermark.size
    width, height = base_image.size

    transparent = Image.new('RGBA', (width+10, height+15), 'white')

    transparent.paste(base_image, (0, 0))
    transparent.paste(watermark, ((width - width1), (height - height1)+10), mask=watermark)

    transparent.save(output_image_path)


async def full_screenshot(driver, save_path):
    # initiate value
    ele=driver.find_element_by_css_selector('body > div:nth-child(2)')
    total_height = ele.size["height"] + 1000
    driver.set_window_size(1920, total_height)
    driver.save_screenshot(save_path)


async def get_option_math_pro(option_number,id)-> list:
    browser.get('https://ege.sdamgia.ru/')
    time.sleep(2.5)
    while (str(browser.current_url) != 'https://ege.sdamgia.ru/'):
        continue
    browser.find_elements_by_css_selector(
        "html.Html body.Page div.Root div.App div.PageLayout.StartPage div.PageLayout-Grid div.PageLayout-Content main.PageLayout-Main section.Section.OurVariants div.OurVariants-Grid a.Link.VariantLink.OurVariants-Link")[
        option_number - 1].click()
    time.sleep(1)
    tasks_urls = []  # получаем ссылки на задания
    tasks = []  # сами задания
    for i in range(2,((19+1)*2),2):
        print(i/2,browser.find_element_by_xpath('/html/body/div[1]/div[5]/div[6]/div['+str(i)+']/div/div[1]/div[1]/span/a').get_attribute('href'))
        tasks_urls.append(str(browser.find_element_by_xpath('/html/body/div[1]/div[5]/div[6]/div['+str(i)+']/div/div[1]/div[1]/span/a').get_attribute('href')))
    for i in range(len(tasks_urls)):

        # скриншот ответа
        browser.get(f"https://ege.sdamgia.ru/problem?id={tasks_urls[i].replace('https://ege.sdamgia.ru/problem?id=','')}&print=true&svg=0&num=true&sol=true&key=true")
        time.sleep(2)
        size = browser.find_element_by_class_name('prob_maindiv').size
        total_height = (size['height'] - 30) #-ответ

        await full_screenshot(browser, 'tasks\\full_' + str(i + 1) + "_" + str(id) + '.png')
        location = browser.find_element_by_xpath(
            '//*[@id="sol' + tasks_urls[i].replace('https://ege.sdamgia.ru/problem?id=','') + '"]').location
        size = browser.find_element_by_xpath(
            '//*[@id="sol' + tasks_urls[i].replace('https://ege.sdamgia.ru/problem?id=','') + '"]').size
        x = location['x'];
        y = location['y'];
        width = location['x'] + size['width']+5;
        heightanswer = size['height']
        height = location['y'] + size['height'];
        #
        # редактирование
        im = Image.open('tasks\\full_' + str(i + 1) + "_" + str(id) + '.png')
        im = im.crop((int(x), int(y), int(width), int(height)))
        im.save('tasks\\full_answer_' + str(i + 1) + "_" + str(id) + '.png')
        await watermark_with_transparency('tasks\\full_answer_' + str(i + 1) + "_" + str(id) + '.png',
                                    'tasks\\full_answer_' + str(i + 1) + "_" + str(id) + '.png',
                                    'watermark.png')
        # div.prob_answer:nth-child(5) > div:nth-child(2)
        #
        if (i<12):
            answer = browser.find_element_by_css_selector('tr.prob_answer:nth-child(2) > td:nth-child(3)').text.strip()




        #
        # скриншот задания
        location = browser.find_element_by_class_name('pbody').location
        size = browser.find_element_by_class_name('pbody').size
        x = location['x'];
        y = location['y'];
        width = location['x'] + size['width'];
        height = y + (total_height - heightanswer);
        #
        # редактирование

        im = Image.open('tasks\\full_' + str(i + 1) + "_" + str(id) + '.png')
        im = im.crop((int(x), int(y), int(width), int(height)))
        im.save('tasks\\task_' + str(i + 1) + "_" + str(id) + '.png')
        await watermark_with_transparency('tasks\\task_' + str(i + 1) + "_" + str(id) + '.png',
                                    'tasks\\task_' + str(i + 1) + "_" + str(id) + '.png', 'watermark.png')
        #
        # добавление путей картинок и ответа в массив,а а также удаление общего скриншота
        if i<12:
            tasks.append([answer, os.path.join('tasks', 'task_' + str(i + 1) + "_" + str(id) + '.png'),
                          os.path.join('tasks', 'full_answer_' + str(i + 1) + "_" + str(id) + '.png')])
        else:
            tasks.append(['None', os.path.join('tasks', 'task_' + str(i + 1) + "_" + str(id) + '.png'),
                          os.path.join('tasks', 'full_answer_' + str(i + 1) + "_" + str(id) + '.png')])
        os.remove(os.path.join("tasks", 'full_' + str(i + 1) + "_" + str(id) + '.png'))


    return tasks

async def get_options_math_pro(numbers_tasks,id)-> list:
    browser.get('https://ege.sdamgia.ru/')
    # здесь можно доработать
    time.sleep(2.5)
    while (str(browser.current_url)!='https://ege.sdamgia.ru/'):
        continue
    ###########
    if numbers_tasks<=12:
        for i in range(1,numbers_tasks+1):#первые 12 заданий
            button= browser.find_element_by_css_selector("div.ConstructorForm-Row:nth-child("+str(i+1)+") > div:nth-child(1) > button:nth-child(3)")
            button.click()
    else:
        for i in range(1,numbers_tasks+1):#первые 12 заданий
            if i==13:
                break
            button= browser.find_element_by_css_selector("div.ConstructorForm-Row:nth-child("+str(i+1)+") > div:nth-child(1) > button:nth-child(3)")
            button.click()

        for i in range(13,numbers_tasks+1):# развернутые вопросы с 13 по 19
            button = browser.find_element_by_css_selector(
                "div.ConstructorForm-Row:nth-child(" + str(i + 2) + ") > div:nth-child(1) > button:nth-child(3)")
            button.click()

    time.sleep(1)
    button=browser.find_element_by_css_selector('.ConstructorForm-SubmitButton')
    button.click()#получаем вариант
    browser.get(str(browser.current_url).replace('True',"False"))
    tasks_urls=[]#получаем ссылки на задания
    tasks=[]#сами задания

    for i in range(2,((numbers_tasks+1)*2),2):
        print(i/2,browser.find_element_by_xpath('/html/body/div[1]/div[5]/div[3]/div['+str(i)+']/div/div[1]/div[1]/span/a').get_attribute('href'))
        tasks_urls.append(str(browser.find_element_by_xpath('/html/body/div[1]/div[5]/div[3]/div['+str(i)+']/div/div[1]/div[1]/span/a').get_attribute('href')))

    #обрабатываем задания
    for i in range(len(tasks_urls)):

        #скриншот ответа
        browser.get(
            f"https://ege.sdamgia.ru/problem?id={tasks_urls[i].replace('https://ege.sdamgia.ru/problem?id=', '')}&print=true&svg=0&num=true&sol=true&key=true")
        time.sleep(2)
        size = browser.find_element_by_class_name('prob_maindiv').size
        total_height = (size['height'] - 30)  # -ответ

        await full_screenshot(browser, 'tasks\\full_' + str(i + 1) + "_" + str(id) + '.png')
        location = browser.find_element_by_xpath(
            '//*[@id="sol' + tasks_urls[i].replace('https://ege.sdamgia.ru/problem?id=', '') + '"]').location
        size = browser.find_element_by_xpath(
            '//*[@id="sol' + tasks_urls[i].replace('https://ege.sdamgia.ru/problem?id=', '') + '"]').size
        x = location['x'];
        y = location['y'];
        width = location['x'] + size['width'] + 5;
        heightanswer = size['height']
        height = location['y'] + size['height'];
        #
        # редактирование
        im = Image.open('tasks\\full_' + str(i + 1) + "_" + str(id) + '.png')
        im = im.crop((int(x), int(y), int(width), int(height)))
        im.save('tasks\\full_answer_' + str(i + 1) + "_" + str(id) + '.png')
        await watermark_with_transparency('tasks\\full_answer_' + str(i + 1) + "_" + str(id) + '.png',
                                    'tasks\\full_answer_' + str(i + 1) + "_" + str(id) + '.png',
                                    'watermark.png')
        # div.prob_answer:nth-child(5) > div:nth-child(2)
        #
        if (i<12):
            answer = browser.find_element_by_css_selector('tr.prob_answer:nth-child(2) > td:nth-child(3)').text.strip()


        #
        # скриншот задания
        location = browser.find_element_by_class_name('pbody').location
        size = browser.find_element_by_class_name('pbody').size
        x = location['x'];
        y = location['y'];
        width = location['x'] + size['width'];
        height = y + (total_height - heightanswer);
        #
        # редактирование

        im = Image.open('tasks\\full_' + str(i + 1) + "_" + str(id) + '.png')
        im = im.crop((int(x), int(y), int(width), int(height)))
        im.save('tasks\\task_' + str(i + 1) + "_" + str(id) + '.png')
        await watermark_with_transparency('tasks\\task_' + str(i + 1) + "_" + str(id) + '.png',
                                    'tasks\\task_' + str(i + 1) + "_" + str(id) + '.png', 'watermark.png')
        #
        # добавление путей картинок и ответа в массив,а а также удаление общего скриншота
        if i < 12:
            tasks.append([answer, os.path.join('tasks', 'task_' + str(i + 1) + "_" + str(id) + '.png'),
                          os.path.join('tasks', 'full_answer_' + str(i + 1) + "_" + str(id) + '.png')])
        else:
            tasks.append(['None', os.path.join('tasks', 'task_' + str(i + 1) + "_" + str(id) + '.png'),
                          os.path.join('tasks', 'full_answer_' + str(i + 1) + "_" + str(id) + '.png')])
        os.remove(os.path.join("tasks", 'full_' + str(i + 1) + "_" + str(id) + '.png'))

    return tasks


async def get_options_math(numbers_tasks,id)-> list:
    browser.get('https://mathb-ege.sdamgia.ru/')
    # здесь можно доработать
    time.sleep(2.5)
    while (str(browser.current_url)!='https://mathb-ege.sdamgia.ru/'):
        continue
    ###########
    
    for i in range(1,numbers_tasks+1):
            button= browser.find_element_by_css_selector("div.ConstructorForm-Row:nth-child("+str(i+1)+") > div:nth-child(1) > button:nth-child(3)")
            button.click()
    

    time.sleep(1)
    button=browser.find_element_by_css_selector('.ConstructorForm-SubmitButton')
    button.click()#получаем вариант
    browser.get(str(browser.current_url).replace('True',"False"))
    tasks_urls=[]#получаем ссылки на задания
    tasks=[]#сами задания

    for i in range(2,((numbers_tasks+1)*2),2):
        print(i/2,browser.find_element_by_xpath('/html/body/div[1]/div[5]/div[3]/div['+str(i)+']/div/div[1]/div[1]/span/a').get_attribute('href'))
        tasks_urls.append(str(browser.find_element_by_xpath('/html/body/div[1]/div[5]/div[3]/div['+str(i)+']/div/div[1]/div[1]/span/a').get_attribute('href')))

    #обрабатываем задания
    for i in range(len(tasks_urls)):

        #скриншот ответа
        browser.get(
            f"https://mathb-ege.sdamgia.ru/problem?id={tasks_urls[i].replace('https://mathb-ege.sdamgia.ru/problem?id=', '')}&print=true&svg=0&num=true&sol=true&key=true")
        time.sleep(2)
        size = browser.find_element_by_class_name('prob_maindiv').size
        total_height = (size['height'] - 30)  # -ответ

        await full_screenshot(browser, 'tasks\\full_' + str(i + 1) + "_" + str(id) + '.png')
        location = browser.find_element_by_xpath(
            '//*[@id="sol' + tasks_urls[i].replace('https://mathb-ege.sdamgia.ru/problem?id=', '') + '"]').location
        size = browser.find_element_by_xpath(
            '//*[@id="sol' + tasks_urls[i].replace('https://mathb-ege.sdamgia.ru/problem?id=', '') + '"]').size
        x = location['x'];
        y = location['y'];
        width = location['x'] + size['width'] + 5;
        heightanswer = size['height']
        height = location['y'] + size['height'];
        #
        # редактирование
        im = Image.open('tasks\\full_' + str(i + 1) + "_" + str(id) + '.png')
        im = im.crop((int(x), int(y), int(width), int(height)))
        im.save('tasks\\full_answer_' + str(i + 1) + "_" + str(id) + '.png')
        await watermark_with_transparency('tasks\\full_answer_' + str(i + 1) + "_" + str(id) + '.png',
                                    'tasks\\full_answer_' + str(i + 1) + "_" + str(id) + '.png',
                                    'watermark.png')
        # div.prob_answer:nth-child(5) > div:nth-child(2)
        #
        answer = browser.find_element_by_css_selector('tr.prob_answer:nth-child(2) > td:nth-child(3)').text.strip()


        #
        # скриншот задания
        location = browser.find_element_by_class_name('pbody').location
        size = browser.find_element_by_class_name('pbody').size
        x = location['x'];
        y = location['y'];
        width = location['x'] + size['width'];
        height = y + (total_height - heightanswer);
        #
        # редактирование

        im = Image.open('tasks\\full_' + str(i + 1) + "_" + str(id) + '.png')
        im = im.crop((int(x), int(y), int(width), int(height)))
        im.save('tasks\\task_' + str(i + 1) + "_" + str(id) + '.png')
        await watermark_with_transparency('tasks\\task_' + str(i + 1) + "_" + str(id) + '.png',
                                    'tasks\\task_' + str(i + 1) + "_" + str(id) + '.png', 'watermark.png')
        #
        # добавление путей картинок и ответа в массив,а а также удаление общего скриншота

        tasks.append([answer, os.path.join('tasks', 'task_' + str(i + 1) + "_" + str(id) + '.png'),
                          os.path.join('tasks', 'full_answer_' + str(i + 1) + "_" + str(id) + '.png')])

        os.remove(os.path.join("tasks", 'full_' + str(i + 1) + "_" + str(id) + '.png'))

    return tasks



async def get_option_math(option_number,id)-> list:
    browser.get('https://mathb-ege.sdamgia.ru/')
    time.sleep(2.5)
    while (str(browser.current_url) != 'https://mathb-ege.sdamgia.ru/'):
        continue
    browser.find_elements_by_css_selector(
        "html.Html body.Page div.Root div.App div.PageLayout.StartPage div.PageLayout-Grid div.PageLayout-Content main.PageLayout-Main section.Section.OurVariants div.OurVariants-Grid a.Link.VariantLink.OurVariants-Link")[
        option_number - 1].click()
    time.sleep(1)
    tasks_urls = []  # получаем ссылки на задания
    tasks = []  # сами задания

   
    for i in range(2,((20+1)*2),2):
        print(i/2,browser.find_element_by_xpath('/html/body/div[1]/div[5]/div[6]/div['+str(i)+']/div/div[1]/div[1]/span/a').get_attribute('href'))
        tasks_urls.append(str(browser.find_element_by_xpath('/html/body/div[1]/div[5]/div[6]/div['+str(i)+']/div/div[1]/div[1]/span/a').get_attribute('href')))
    for i in range(len(tasks_urls)):
        
        # скриншот ответа
        browser.get(f"https://mathb-ege.sdamgia.ru/problem?id={tasks_urls[i].replace('https://mathb-ege.sdamgia.ru/problem?id=','')}&print=true&svg=0&num=true&sol=true&key=true")
        time.sleep(2)
        size = browser.find_element_by_class_name('prob_maindiv').size
        total_height = (size['height'] - 30) #-ответ

        await full_screenshot(browser, 'tasks\\full_' + str(i + 1) + "_" + str(id) + '.png')
        location = browser.find_element_by_xpath(
            '//*[@id="sol' + tasks_urls[i].replace('https://mathb-ege.sdamgia.ru/problem?id=','') + '"]').location
        size = browser.find_element_by_xpath(
            '//*[@id="sol' + tasks_urls[i].replace('https://mathb-ege.sdamgia.ru/problem?id=','') + '"]').size
        x = location['x'];
        y = location['y'];
        width = location['x'] + size['width']+5;
        heightanswer = size['height']
        height = location['y'] + size['height'];
        #
        # редактирование
        im = Image.open('tasks\\full_' + str(i + 1) + "_" + str(id) + '.png')
        im = im.crop((int(x), int(y), int(width), int(height)))
        im.save('tasks\\full_answer_' + str(i + 1) + "_" + str(id) + '.png')
        await watermark_with_transparency('tasks\\full_answer_' + str(i + 1) + "_" + str(id) + '.png',
                                    'tasks\\full_answer_' + str(i + 1) + "_" + str(id) + '.png',
                                    'watermark.png')
        # div.prob_answer:nth-child(5) > div:nth-child(2)
        #tr.prob_answer:nth-child(2) > td:nth-child(3)
        answer = browser.find_element_by_css_selector('tr.prob_answer:nth-child(2) > td:nth-child(3)').text.strip()




        #
        # скриншот задания
        location = browser.find_element_by_class_name('pbody').location
        size = browser.find_element_by_class_name('pbody').size
        x = location['x'];
        y = location['y'];
        width = location['x'] + size['width'];
        height = y + (total_height - heightanswer);
        #
        # редактирование

        im = Image.open('tasks\\full_' + str(i + 1) + "_" + str(id) + '.png')
        im = im.crop((int(x), int(y), int(width), int(height)))
        im.save('tasks\\task_' + str(i + 1) + "_" + str(id) + '.png')
        await watermark_with_transparency('tasks\\task_' + str(i + 1) + "_" + str(id) + '.png',
                                    'tasks\\task_' + str(i + 1) + "_" + str(id) + '.png', 'watermark.png')
        #
        # добавление путей картинок и ответа в массив,а а также удаление общего скриншота
        tasks.append([answer, os.path.join('tasks', 'task_' + str(i + 1) + "_" + str(id) + '.png'),
                          os.path.join('tasks', 'full_answer_' + str(i + 1) + "_" + str(id) + '.png')])
       
        os.remove(os.path.join("tasks", 'full_' + str(i + 1) + "_" + str(id) + '.png'))


    return tasks


